"""
#===============================================================================
# License
#===============================================================================
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

#===============================================================================
Copyright (c) 2009 bryan p. conrad
bconrad2@gmail.com

Repository Path:    
$$HeadURL$$

Created on:           7/27/2009
Last Revised on:      $$Date$$
Version:              $$Revision$$
ID:                   $$Id$$
Author:               $$Author$$ :: (bryan)

#===============================================================================


#===============================================================================
# Documentation
#===============================================================================

"""

#===============================================================================
# Imports
#===============================================================================
import JT
from JT.Optimizer import RegularStepGradientDescentOptimizer
import numpy as np
import my_auth
import time
import csv
import os
import sys
from socket import gethostname
from datetime import datetime
import time


class Test(object):
    """
    this module will make it easy to construct registration tests. 
    """

    def __init__(self):
        self._connectSuccess = False
        self._uploadData = False
        self._logData = False
        
    def setup(self,upload=True,logging=True):
#===============================================================================
#    Google Data Upload Setup
#===============================================================================
        cred = my_auth.credentials()
        username = cred.name
        password = cred.pw
        
        if upload:
            self._uploadData = True
            try:
                self.uploader = JT.Uploader(username,password)
                time.sleep(1)
                # Choose newest sheet and book
                self.uploader._run()
                self._connectSuccess = True
            except:
                print "Unexpected error:", sys.exc_info()[0]
                print 'Could not connect to Google'
    
        if logging:
            self._logData = True
            logTag = 'Phase2'     #This tag can be used to identify trials
            logName = time.strftime('%Y%m%d_%H%M%S')+'_'+logTag+'.log'
            logDir = 'C:/Users/bryan/bryan-code/trunk/Results'
            print 'Final Result Logging dir: ', logDir 
            self.logFile = open(os.path.join(logDir,logName),'wt')
            self.log = csv.writer(self.logFile,dialect='excel-tab')
            # Write Headers
            self.log.writerow(('user','time','n_pose', 'method', 'final_cost', 
                               'vert', 'proj', 'body',
                               'Starting Pose0(mm)',
                               'Starting Pose1(mm)','Starting Pose2(mm)',
                               'Starting Pose3(deg)','Starting Pose4(deg)',
                               'Starting Pose5(deg)','Final Pose0(mm)',
                               'Final Pose1(mm)','Final Pose2(mm)',
                               'Final Pose3(deg)','Final Pose4(deg)',
                               'Final Pose5(deg)','mTRE','mTRE_proj'))
     
        
        # First setup the JointTrack data structure
        self.data = JT.Data2.gold_standard()     
#===============================================================================
# Setup Registration
#===============================================================================
        self.reg=JT.MultiIterationRegistration()
        # Set additional arguments for cost function
        
        self.reg.settings['plot'] = False
        self.reg.settings['debug'] = True
        self.reg._logGradientResult = False
        self.reg.SetIterations(1)

#===============================================================================
#    Setup Optimizer
#===============================================================================      
        gd = RegularStepGradientDescentOptimizer()
        self.reg.SetOptimizer(gd)
        # Set the cost function
        gd.SetCostFunction(self.reg)
        # Modify Optimizer Parameters
        self.reg.SetInitialStepLength(.5)
        gd.SetMinimumStepLength(.1)
        gd.SetRelaxationFactor(0.9)
        #gd.SetMaximumGradientStepLength(0.5)
        gd.SetNumberOfIterations(50)
        gd._debug = True
    
        
    def run(self,n_pose=50,method='bpc',vert=0,proj=0,body=0):   
#===============================================================================
# Setup data 
#===============================================================================
        self.reg.SetFixedImageFileName(str(self.data.Tree.Frame[vert].FixedImage[0].path))
        #self.reg.SetRegionOfInterest(self.data.Frames[vert].FixedImages[proj].ROIs[body].data)
        drr = JT.DRR()
        self.reg.SetDRR(drr)
        original_transform = np.asarray([[0.9965045488,    0.0280853197,    0.0786759118,   -54.388],
                                        [-0.033833201,    0.9967795231,    0.0727041732,   -66.8634],
                                        [-0.0763806179,   -0.0751118973,    0.9942455452,   -43.9706],
                                        [0.0000000000,    0.0000000000,    0.0000000000,    1.0000000000]])
        self.reg.drr.volume.SetOriginalTransform(original_transform)
        #self.reg.drr.renWin.SetOffScreenRendering(0)
        self.reg.drr.OffScreenRenderingOff()
        self.reg.drr.LoadCalFile(str(self.data.Tree.Frame[vert].FixedImage[0].Cal.path))
        self.reg.drr.xrayCamera.DoCalibration()
        self.reg.SetVolumeFileName(str(self.data.Tree.Frame[vert].Volume[0].path))
        self.reg.SetRegionOfInterest(self.data.roi[vert][proj][body])
        self.reg.SetCenter(self.data.centers[vert][body])
        self.reg.SetImageMetric(method)
        
        #print self.reg.drr.volume
        #print self.reg.drr.xrayCamera
        
        # To calculate the error in the projection direction, we need to know
        # the normal to the projection plane which will vary for each 
        # vertebrae and each projection. 
        if (vert is 0) and (proj is 0):
            # vert 1, AP projection
            proj_norm = [0.382335,  0.0106653, -0.923962]
        elif (vert is 0) and (proj is 1):
            # vert 1, LAT projection
            proj_norm = [-0.928639,-0.0257142, -0.370091]
        elif (vert is 1) and (proj is 0):
            # vert 2, AP projection
            proj_norm = [-0.834327, -0.0236362, -0.550762]
        elif (vert is 1) and (proj is 1):
            # vert 2, LAT projection
            proj_norm = [0.456834, 0.0115521, -0.889477]
        
        # TODO: The following two calculations seem to be useful and should 
        # probably be moved into the registration class so that they can be 
        # reused by other codes.
        # adjust scales based on projection direction
#        trans_scales = np.ones(3)*1.1 - abs(np.asarray(proj_norm)) 
#        scales = [trans_scales[0],trans_scales[1],trans_scales[2],
#                  1900.0,1900.0,1900.0]
        scales = [1,1,1,1000,1000,1000]
        gd = self.reg.GetOptimizer()
        gd.SetScales(scales)
        
        # Adjust gradient step size to take a bigger step in the proj direction
        #trans_grad = np.ones(3)*1.0 + abs(np.asarray(proj_norm)) 
        #grad_step = [trans_grad[0],trans_grad[1],trans_grad[2],0.75,0.75,0.75]
        grad_step = [.25,.25,.25,.1,.1,.1]
        gd.SetMaximumGradientStepLength(grad_step)
        
        # Set the starting pose
        Testing = JT.Testing()
        pose = Testing.GetStartPose(n_pose)
        self.reg.SetInitialPose(pose)
        self.reg._initialStartingPose = pose
        self.reg.metric_settings['gsigma'] = 2
        self.reg.metric_settings['debug'] = False
        
        #self.drr._volume_list[0].UseRayCastMapper()
        print "Starting Pose: ", pose
        self.reg.Start()

        # Save registration data to file for analysis of tuning parameters
        # Log status/results to a file that can be plotted later.
        initial_pose = pose
        final_pose = self.reg.bestPose
        final_cost = self.reg.bestCost
        
        try:
            [mTRE, mTRE_proj] = JT.utils.FindMTre(final_pose,proj_norm)
        except:
            print "Error calculating mTRE"
            [mTRE,mTRE_proj] = 9999,9999
        dt = datetime.fromtimestamp(time.time())
        time_stamp = dt.strftime('%m/%d/%Y_%H:%M:%S')
        user = gethostname()
        
        my_tupel = [user,time_stamp,n_pose, method, final_cost, vert, 
                    proj, body, initial_pose[0], initial_pose[1],
                    initial_pose[2],np.degrees(initial_pose[3]), 
                    np.degrees(initial_pose[4]),np.degrees(initial_pose[5]),
                    final_pose[0],final_pose[1],final_pose[2],
                    np.degrees(final_pose[3]),np.degrees(final_pose[4]),
                    np.degrees(final_pose[5]),mTRE,mTRE_proj]
        
        if self._logData:
            self.log.writerow(my_tupel)
            self.logFile.flush()
        
        if self._uploadData and self._connectSuccess:
            output_dict = {}
            output_dict['user'] = user
            output_dict['time'] = time_stamp
            output_dict['npose'] = str(n_pose)
            output_dict['method'] = method
            output_dict['cost'] = str(final_cost)
            output_dict['vert'] = str(vert)
            output_dict['projection'] = str(proj)
            output_dict['body'] = str(body)
            output_dict['startingpose0'] = str(initial_pose[0])
            output_dict['startingpose1'] = str(initial_pose[1])
            output_dict['startingpose2'] = str(initial_pose[2])
            output_dict['startingpose3'] = str(np.degrees(initial_pose[3]))
            output_dict['startingpose4'] = str(np.degrees(initial_pose[4]))
            output_dict['startingpose5'] = str(np.degrees(initial_pose[5]))
            output_dict['finalpose0'] = str(final_pose[0])
            output_dict['finalpose1'] = str(final_pose[1])
            output_dict['finalpose2'] = str(final_pose[2])
            output_dict['finalpose3'] = str(np.degrees(final_pose[3]))
            output_dict['finalpose4'] = str(np.degrees(final_pose[4]))
            output_dict['finalpose5'] = str(np.degrees(final_pose[5]))
            output_dict['stepsize'] = str(self.reg._initialStepLength)
            output_dict['mtre'] = str(mTRE)
            output_dict['mtrep'] = str(mTRE_proj)
            self.uploader.Write(output_dict)
    

if __name__ == "__main__":
    
    test = Test()
    test.setup(upload=False, logging=True)
    user = gethostname()
    #method_list = ['nmi', 'mi', 'ncc', 'ecc', 'mse', 'ncc', 'gd', 'gc']
    #method_list = ['bpc']
    if user == 'Glenoid':
        method_list = ['gc','bpc','ncc','gd','mse']
    else:
        method_list = ['mi','nmi','ecc','mse']
                              
    vert = [0,1]
    proj = [1,0]
    body = [[0,1,2],[0,1,2,3,4]]
    #body=[[1],[3]]
    if 1:
        v=0
        p=0
        b=1
        n_pose=14
        method = 'ncc'
        tag = 'Method: %s, Vert: %s, Proj: %s, Body: %s, Pose: %i' %\
              (method,str(v),str(p),str(b),n_pose)
        print tag
        test.run(n_pose=n_pose,method=method,vert=v,proj=p,body=b)
        test.reg.ShowFinalResults()
    else:
        for method in method_list:
            for p in proj:
                for v in vert:
                    for b in body[v]:
                        for n_pose in range(150):               
                            # to label the resulting image that will be saved.
                            tag = 'Method: %s, Vert: %s, Proj: %s, Body: %s, Pose: %i' %\
                              (method,str(v),str(p),str(b),n_pose)
                            print tag
                            test.run(n_pose=n_pose,method=method,
                                     vert=v,proj=p,body=b)
                        
                        
                        